package com.lnn.springBase.utils.sftp;

import com.jcraft.jsch.ChannelSftp;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.Session;
import com.jcraft.jsch.SftpException;

import java.io.FileOutputStream;
import java.io.OutputStream;
import java.util.Properties;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class SftpUtil {

    private String host;
    private int port;
    private String username;
    private String password;
    private ChannelSftp sftpChannel;
    private Session session;

    public SftpUtil(String host, int port, String username, String password) {
        this.host = host;
        this.port = port;
        this.username = username;
        this.password = password;
    }

    public void connect() throws Exception {
        JSch jsch = new JSch();
        session = jsch.getSession(username, host, port);
        session.setPassword(password);
        Properties config = new Properties();
        config.put("StrictHostKeyChecking", "no");
        session.setConfig(config);
        session.connect();
        sftpChannel = (ChannelSftp) session.openChannel("sftp");
        sftpChannel.connect();
    }

    public void disconnect() {
        if (sftpChannel != null && sftpChannel.isConnected()) {
            sftpChannel.disconnect();
        }
        if (session != null && session.isConnected()) {
            session.disconnect();
        }
    }

    public void downloadFile(String remoteFile, String localFile) throws Exception {
        try {
            downloadFileSingleThread(remoteFile, localFile);
        } finally {
            disconnect();
        }
    }

    public void downloadFile(String remoteFile, String localFile,int chunkSize, int threadCount) throws Exception {
        try {
            downloadFileInChunks(remoteFile, localFile, chunkSize, threadCount);
        } finally {
            disconnect();
        }
    }

    private void downloadFileSingleThread(String remoteFile, String localFile) throws SftpException {
        ProgressMonitor monitor = new ProgressMonitor(sftpChannel.lstat(remoteFile).getSize());
        try (OutputStream outputStream = new FileOutputStream(localFile)) {
            sftpChannel.get(remoteFile, outputStream, monitor);
        } catch (Exception e) {
            throw new SftpException(ChannelSftp.SSH_FX_FAILURE, "Failed to download file", e);
        }
    }

    private void downloadFileInChunks(String remoteFile, String localFile, int chunkSize, int threadCount) throws Exception {
        long fileSize = sftpChannel.lstat(remoteFile).getSize();
        ProgressMonitor monitor = new ProgressMonitor(fileSize);
        ExecutorService executor = Executors.newFixedThreadPool(threadCount);

        for (long offset = 0; offset < fileSize; offset += chunkSize) {
            long end = Math.min(offset + chunkSize, fileSize);
            executor.submit(new SFTPDownloadTask(sftpChannel, remoteFile, localFile, offset, end, monitor));
        }

        executor.shutdown();
        executor.awaitTermination(Long.MAX_VALUE, TimeUnit.NANOSECONDS);
    }

    private static class SFTPDownloadTask implements Runnable {
        private final ChannelSftp sftpChannel;
        private final String remoteFile;
        private final String localFile;
        private final long start;
        private final long end;
        private final ProgressMonitor monitor;

        public SFTPDownloadTask(ChannelSftp sftpChannel, String remoteFile, String localFile, long start, long end, ProgressMonitor monitor) {
            this.sftpChannel = sftpChannel;
            this.remoteFile = remoteFile;
            this.localFile = localFile;
            this.start = start;
            this.end = end;
            this.monitor = monitor;
        }

        @Override
        public void run() {
            try (OutputStream outputStream = new FileOutputStream(localFile + "." + start + "-" + end)) {
                sftpChannel.get(remoteFile, outputStream, monitor, ChannelSftp.RESUME, start);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private static class ProgressMonitor implements com.jcraft.jsch.SftpProgressMonitor {
        private long totalSize;
        private long downloadedSize;

        public ProgressMonitor(long totalSize) {
            this.totalSize = totalSize;
            this.downloadedSize = 0;
        }

        @Override
        public void init(int op, String src, String dest, long max) {
            System.out.println("Starting download: " + src);
        }

        @Override
        public boolean count(long count) {
            downloadedSize += count;
            int progress = (int) ((downloadedSize * 100) / totalSize);
            System.out.print("\rDownloaded: " + progress + "%");
            return true;
        }

        @Override
        public void end() {
            System.out.println("\nDownload completed.");
        }
    }

    public static void main(String[] args) {
        String host = "211.144.120.164";
        int port = 60022;
        String username = "sftpuser";
        String password = "yxkj2025";
        SftpUtil sftpUtil = new SftpUtil(host, port, username, password);
        try {
            sftpUtil.connect();// 多线程时使用的线程数
            sftpUtil.downloadFile("/data/sftp/EC_2024121200.nc.tar.gz", "E:\\工作\\FD1\\EC_2024121200.nc.tar.gz"); // 每个块1MB
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            sftpUtil.disconnect();
        }
    }
}
